#!/bin/bash
#SBATCH --job-name=nsem_scaling
#SBATCH --account=your_project_account
#SBATCH --partition=regular
#SBATCH --time=02:00:00
#SBATCH --nodes=16
#SBATCH --ntasks-per-node=32
#SBATCH --cpus-per-task=1
#SBATCH --constraint=haswell
#SBATCH --output=nsem_scaling_%j.out
#SBATCH --error=nsem_scaling_%j.err
#SBATCH --mail-type=ALL
#SBATCH --mail-user=your_email@institution.edu

# NSEMSolver.jl HPC Scaling Study
# This script runs strong and weak scaling studies for NSEMSolver.jl

echo "Starting NSEMSolver.jl HPC Scaling Study"
echo "========================================"
echo "Job ID: $SLURM_JOB_ID"
echo "Nodes: $SLURM_JOB_NUM_NODES"
echo "Tasks: $SLURM_NTASKS"
echo "Start time: $(date)"
echo

# Load required modules
module purge
module load julia/1.9.0
module load openmpi/4.1.0
module load mkl/2021.4.0

# Set environment variables for optimal performance
export JULIA_NUM_THREADS=1
export OMP_NUM_THREADS=1
export JULIA_MPI_BINARY=system
export JULIA_CUDA_USE_BINARYBUILDER=false

# Julia environment setup
export JULIA_DEPOT_PATH=$SCRATCH/.julia
export JULIA_PROJECT=$SLURM_SUBMIT_DIR

echo "Environment Setup:"
echo "Julia version: $(julia --version)"
echo "MPI version: $(mpirun --version | head -1)"
echo "Julia threads: $JULIA_NUM_THREADS"
echo "Julia project: $JULIA_PROJECT"
echo

# Initialize Julia environment
echo "Initializing Julia environment..."
julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()'

# Define scaling study parameters
PROBLEM_SIZES=(16 24 32 48 64)
PROCESSOR_COUNTS=(32 64 128 256 512)

# Create output directory
OUTPUT_DIR="scaling_results_$(date +%Y%m%d_%H%M%S)"
mkdir -p $OUTPUT_DIR

echo "Starting Strong Scaling Study"
echo "============================="

# Strong scaling: fixed problem size, varying processor count
FIXED_PROBLEM_SIZE=64
for NPROCS in "${PROCESSOR_COUNTS[@]}"; do
    if [ $NPROCS -le $SLURM_NTASKS ]; then
        echo "Running strong scaling test with $NPROCS processors..."
        
        # Calculate nodes needed
        NODES_NEEDED=$(( (NPROCS + 31) / 32 ))
        if [ $NODES_NEEDED -le $SLURM_JOB_NUM_NODES ]; then
            
            OUTPUT_FILE="$OUTPUT_DIR/strong_scaling_${NPROCS}procs.out"
            
            srun --nodes=$NODES_NEEDED --ntasks=$NPROCS --ntasks-per-node=32 \
                 julia --project=. examples/hpc/hpc_scaling_demo.jl \
                 --problem-size=$FIXED_PROBLEM_SIZE \
                 --study-type=strong \
                 --output-file=$OUTPUT_FILE
            
            echo "Completed strong scaling test with $NPROCS processors"
        else
            echo "Skipping $NPROCS processors (exceeds allocated nodes)"
        fi
    fi
done

echo
echo "Starting Weak Scaling Study"
echo "=========================="

# Weak scaling: proportional problem size and processor count
for i in "${!PROCESSOR_COUNTS[@]}"; do
    NPROCS=${PROCESSOR_COUNTS[$i]}
    PROBLEM_SIZE=${PROBLEM_SIZES[$i]}
    
    if [ $NPROCS -le $SLURM_NTASKS ]; then
        echo "Running weak scaling test with $NPROCS processors, problem size $PROBLEM_SIZE..."
        
        # Calculate nodes needed
        NODES_NEEDED=$(( (NPROCS + 31) / 32 ))
        if [ $NODES_NEEDED -le $SLURM_JOB_NUM_NODES ]; then
            
            OUTPUT_FILE="$OUTPUT_DIR/weak_scaling_${NPROCS}procs_${PROBLEM_SIZE}size.out"
            
            srun --nodes=$NODES_NEEDED --ntasks=$NPROCS --ntasks-per-node=32 \
                 julia --project=. examples/hpc/hpc_scaling_demo.jl \
                 --problem-size=$PROBLEM_SIZE \
                 --study-type=weak \
                 --output-file=$OUTPUT_FILE
            
            echo "Completed weak scaling test with $NPROCS processors"
        else
            echo "Skipping $NPROCS processors (exceeds allocated nodes)"
        fi
    fi
done

echo
echo "Starting Thermal Convection Study"
echo "================================"

# Large-scale thermal convection simulation
THERMAL_PROCS=256
if [ $THERMAL_PROCS -le $SLURM_NTASKS ]; then
    echo "Running large-scale thermal convection with $THERMAL_PROCS processors..."
    
    THERMAL_NODES=$(( (THERMAL_PROCS + 31) / 32 ))
    if [ $THERMAL_NODES -le $SLURM_JOB_NUM_NODES ]; then
        
        OUTPUT_FILE="$OUTPUT_DIR/thermal_convection_${THERMAL_PROCS}procs.out"
        
        srun --nodes=$THERMAL_NODES --ntasks=$THERMAL_PROCS --ntasks-per-node=32 \
             julia --project=. examples/hpc/parallel_thermal_convection.jl \
             --ra-number=1e6 \
             --output-file=$OUTPUT_FILE
        
        echo "Completed thermal convection simulation"
    else
        echo "Skipping thermal convection (exceeds allocated nodes)"
    fi
fi

echo
echo "Starting Turbulent Flow Study"
echo "============================"

# Large-scale turbulent flow simulation
TURBULENT_PROCS=512
if [ $TURBULENT_PROCS -le $SLURM_NTASKS ]; then
    echo "Running large-scale turbulent flow with $TURBULENT_PROCS processors..."
    
    TURBULENT_NODES=$(( (TURBULENT_PROCS + 31) / 32 ))
    if [ $TURBULENT_NODES -le $SLURM_JOB_NUM_NODES ]; then
        
        OUTPUT_FILE="$OUTPUT_DIR/turbulent_flow_${TURBULENT_PROCS}procs.out"
        
        srun --nodes=$TURBULENT_NODES --ntasks=$TURBULENT_PROCS --ntasks-per-node=32 \
             julia --project=. examples/hpc/distributed_turbulent_flow.jl \
             --reynolds-number=5000 \
             --output-file=$OUTPUT_FILE
        
        echo "Completed turbulent flow simulation"
    else
        echo "Skipping turbulent flow (exceeds allocated nodes)"
    fi
fi

# Generate performance report
echo
echo "Generating Performance Analysis Report"
echo "===================================="

julia --project=. -e "
using NSEMSolver

# Analyze all results in output directory
output_dir = \"$OUTPUT_DIR\"
study = ScalingStudy(\"HPC Scaling Study\", \"Comprehensive scaling analysis\")

# Load and analyze results (would parse output files)
# This is a placeholder - actual implementation would parse the output files

generate_performance_report(study, joinpath(output_dir, \"scaling_analysis_report.md\"))
"

echo
echo "Scaling Study Completed Successfully"
echo "==================================="
echo "End time: $(date)"
echo "Output directory: $OUTPUT_DIR"
echo "View results with: cat $OUTPUT_DIR/scaling_analysis_report.md"
echo

# Archive results
tar -czf "${OUTPUT_DIR}.tar.gz" $OUTPUT_DIR
echo "Results archived to: ${OUTPUT_DIR}.tar.gz"